import pprint  # 较美观的列印函数
import gensim
from collections import defaultdict
from gensim import corpora


documents = [
    "Human machine interface for lab abc computer applications",
    "A survey of user opinion of computer system response time",
    "The EPS user interface management system",
    "System and human system engineering testing of EPS",
    "Relation of user perceived response time to error measurement",
    "The generation of random binary unordered trees",
    "The intersection graph of paths in trees",
    "Graph minors IV Widths of trees and well quasi ordering",
    "Graph minors A survey",
]

# 任意设定一些停用词
stoplist = set('for a of the and to in'.split())

# 分词，转小写
texts = [
    [word for word in document.lower().split() if word not in stoplist]
    for document in documents
]
# print(texts)

frequency = defaultdict(int)
for text in texts:
    for token in text:
        frequency[token] += 1

texts = [
    [token for token in text if frequency[token] > 1]
    for text in texts
]
# 转为字典
dictionary = corpora.Dictionary(texts)

# 转为 BOW
corpus = [dictionary.doc2bow(text) for text in texts]
print(corpus)

# 建立 LSI (Latent semantic indexing) 模型
from gensim import models

# num_topics=2：取二维，即两个议题
lsi = models.LsiModel(corpus, id2word=dictionary, num_topics=2)

# 两个议题的 LSI 公式
lsi.print_topics(2)

# 例句
doc = "Human computer interaction"

# 测试 LSI (Latent semantic indexing) 模型
vec_bow = dictionary.doc2bow(doc.lower().split())
vec_lsi = lsi[vec_bow]
print(vec_lsi)

from gensim import similarities

# 比较例句与语料库的相似性索引
index = similarities.MatrixSimilarity(lsi[corpus])

# 比较例句与语料库的相似机率
sims = index[vec_lsi]

# 显示语料库的索引值及相似机率
print(list(enumerate(sims)))

sims = sorted(enumerate(sims), key=lambda item: -item[1])
for doc_position, doc_score in sims:
    print(doc_score, documents[doc_position])
